package com.loongson;

import com.sun.jna.Memory;
import com.sun.jna.Pointer;
import com.sun.jna.platform.win32.*;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.util.Base64;

public class AppIcon {

    public static String extractIcon(String exePath) throws IOException {
        // 获取图标句柄
        WinDef.HICON[] largeIcon = new WinDef.HICON[1];
        WinDef.HICON[] smallIcon = new WinDef.HICON[1];

        // ExtractIconEx 需要 int 类型的参数
        int iconCount = Shell32.INSTANCE.ExtractIconEx(exePath, 0, largeIcon, smallIcon, 1);
        if (iconCount > 0 && largeIcon[0] != null) {
            // 提取图标信息
            WinGDI.ICONINFO iconInfo = new WinGDI.ICONINFO();
            boolean success = User32.INSTANCE.GetIconInfo(largeIcon[0], iconInfo);
            if (success) {
                BufferedImage iconImage = getBufferedImageFromIcon(iconInfo);
                if (iconImage != null) {
                    // 将图像转换为Base64
                    return bufferedImageToBase64(iconImage);
//                    ImageIO.write(iconImage, "png", new File(outputPath));
//                    System.out.println("Icon saved to " + outputPath);
                }
            }
        } else {
            System.out.println("No icon found in the specified file.");
            String base64IconCmd = extractIcon("C:\\Windows\\System32\\cmd.exe");
            return base64IconCmd;
        }

        return null;
    }

    private static BufferedImage getBufferedImageFromIcon(WinGDI.ICONINFO iconInfo) {
        WinGDI.BITMAP bmp = new WinGDI.BITMAP();
        GDI32.INSTANCE.GetObject(iconInfo.hbmColor, bmp.size(), bmp.getPointer());
        bmp.read();

        int width = bmp.bmWidth.intValue(); // 转换 NativeLong 为 int
        int height = Math.abs(bmp.bmHeight.intValue()); // 转换 NativeLong 为 int

        // Create a bitmap info structure
        WinGDI.BITMAPINFO bmi = new WinGDI.BITMAPINFO();
        bmi.bmiHeader.biSize = bmi.bmiHeader.size();
        bmi.bmiHeader.biWidth = width;
        bmi.bmiHeader.biHeight = -height; // Top-down bitmap
        bmi.bmiHeader.biPlanes = 1;
        bmi.bmiHeader.biBitCount = 32;
        bmi.bmiHeader.biCompression = WinGDI.BI_RGB;

        // Allocate memory for pixel data
        Pointer bits = new Memory(width * height * 4); // 4 bytes per pixel for ARGB
        GDI32.INSTANCE.GetDIBits(User32.INSTANCE.GetDC(null), iconInfo.hbmColor, 0, height, bits, bmi, WinGDI.DIB_RGB_COLORS);

        // Read pixel data into an array
        int[] pixels = bits.getIntArray(0, width * height);

        // Create a BufferedImage
        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB);
        image.setRGB(0, 0, width, height, pixels, 0, width);

        return image;
    }

    private static String bufferedImageToBase64(BufferedImage image) throws IOException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        ImageIO.write(image, "png", baos);
        byte[] imageBytes = baos.toByteArray();
        baos.close();

        return Base64.getEncoder().encodeToString(imageBytes);
    }
}